import torch
from models.net import LeNet
from torchvision.transforms import Compose, ToTensor, Grayscale, Resize
from PIL import Image
import matplotlib.pyplot as plt

# 定义图像预处理
data_transform = Compose([
    Grayscale(),  # 转为灰度图
    Resize((28, 28)),  # 调整为28x28像素
    ToTensor()  # 转为Tensor格式
])

# 如果有NVIDIA显卡，转到GPU训练，否则用CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 模型实例化，将模型转到device
model = LeNet().to(device)

# 加载训练好的模型
model.load_state_dict(torch.load("save_model/best_model.pth"))

# 结果类型
classes = [
    "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
]

# 加载自定义手写图片
image_path = "Tmage/0/0_1.png"  # 替换为你手写图片的路径
image = Image.open(image_path)  # 打开图片

# 预处理图片
image = data_transform(image)
image = image.unsqueeze(0).to(device)  # 增加批次维度并转到设备

# 进入验证阶段
model.eval()  # 切换模型为评估模式
with torch.no_grad():
    pred = model(image)
    predicted = classes[torch.argmax(pred[0])]  # 获取预测类别
    print(f'Predicted: "{predicted}"')

# 显示预处理后的手写图片并标记预测结果
plt.figure("Processed Image")
processed_image = image.cpu().squeeze(0).squeeze(0).numpy()  # 压缩成 (28, 28)
plt.imshow(processed_image, cmap="gray")
plt.title(f"Predicted: {predicted}")  # 标记预测结果
plt.axis("off")
plt.show()